// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/functions/prestosql/function_base_test.h>

using kumo::pollux::test::assertEqualVectors;

namespace kumo::pollux::functions::test {

namespace {

class ArrayFrequencyTest : public functions::test::FunctionBaseTest {
 protected:
  void testArrayFrequency(const VectorPtr& expected, const VectorPtr& input) {
    auto result =
        evaluate<BaseVector>("array_frequency(C0)", make_row_vector({input}));
    assertEqualVectors(expected, result);
  }
};

} // namespace

TEST_F(ArrayFrequencyTest, integerArray) {
  auto array = make_nullable_array_vector<int64_t>(
      {{2, 1, 1, -2},
       {},
       {1, 2, 1, 1, 1, 1},
       {-1, std::nullopt, -1, -1},
       {std::numeric_limits<int64_t>::max(),
        std::numeric_limits<int64_t>::max(),
        1,
        std::nullopt,
        0,
        1,
        std::nullopt,
        0}});

  auto expected = make_map_vector<int64_t, int>(
      {{{1, 2}, {2, 1}, {-2, 1}},
       {},
       {{1, 5}, {2, 1}},
       {{-1, 3}},
       {{std::numeric_limits<int64_t>::max(), 2}, {1, 2}, {0, 2}}});

  testArrayFrequency(expected, array);
}

TEST_F(ArrayFrequencyTest, integerArrayWithoutNull) {
  auto array =
      make_array_vector<int64_t>({{2, 1, 1, -2}, {}, {1, 2, 1, 1, 1, 1}});

  auto expected = make_map_vector<int64_t, int>(
      {{{1, 2}, {2, 1}, {-2, 1}}, {}, {{1, 5}, {2, 1}}});

  testArrayFrequency(expected, array);
}

TEST_F(ArrayFrequencyTest, varcharArray) {
  auto array = make_nullable_array_vector<StringView>({
      {"hello"_sv, "world"_sv, "!"_sv, "!"_sv, "!"_sv},
      {},
      {"hello"_sv, "world"_sv, std::nullopt, "!"_sv, "!"_sv},
      {"helloworldhelloworld"_sv,
       "helloworldhelloworld"_sv,
       std::nullopt,
       "!"_sv,
       "!"_sv},
  });

  auto expected = make_map_vector<StringView, int>(
      {{{"hello"_sv, 1}, {"world"_sv, 1}, {"!"_sv, 3}},
       {},
       {{"hello"_sv, 1}, {"world"_sv, 1}, {"!"_sv, 2}},
       {{"helloworldhelloworld"_sv, 2}, {"!"_sv, 2}}});

  testArrayFrequency(expected, array);
}

TEST_F(ArrayFrequencyTest, varcharArrayWithoutNull) {
  auto array = make_nullable_array_vector<StringView>({
      {"hello"_sv, "world"_sv, "!"_sv, "!"_sv, "!"_sv},
      {},
      {"helloworldhelloworld"_sv, "helloworldhelloworld"_sv, "!"_sv, "!"_sv},
  });
  auto expected = make_map_vector<StringView, int>(
      {{{"hello"_sv, 1}, {"world"_sv, 1}, {"!"_sv, 3}},
       {},
       {{"helloworldhelloworld"_sv, 2}, {"!"_sv, 2}}});

  testArrayFrequency(expected, array);
}

} // namespace kumo::pollux::functions::test
